# !/usr/bin/python3
# coding: utf-8

import os, sys
os.environ["CUDA_VISIBLE_DEVICES"] = '0'

import pickle
import time
import random  
import string  

import torch
import argparse
import numpy as np
import json

from collections import Counter
from rivagan import RivaGAN

parser = argparse.ArgumentParser()
parser.add_argument('--root_dir', type=str, default='./data',
                    help='the path to the project root directory.')
parser.add_argument('--checkpoint', type=str, default='./checkpoints/model_h264_epoch300.pt',
                    help='the path of load pre-trained model weights')
parser.add_argument('--watermark_path', type=str, default='./data/water_marking/water_marking_32.pickle',
                    help='the path of 32-bits water marking')
parser.add_argument('--input_videos_dir', type=str, default='./data/input_videos',
                    help='the path of input videos')
parser.add_argument('--output_dir', type=str, default='./data/results_test',
                    help='the path of output results after add watermark messages')
parser.add_argument('--label_path', type=str, default=' ',
                    help='the path of videomark encode label')
parser.add_argument('--threshold', type=float, default=0.15,
                    help='the threshold to distinguish has/no watermark image frame')
parser.add_argument('--use_dictfile', action='store_true', default=False,
                    help='if true, use generated decoded watermarks file last time')
parser.add_argument('--no_accuracy', action='store_true', default=False,
                    help='if compute decoded accuracy')
parser.add_argument('--encode', action='store_true', default=False,
                    help='if use encoder to embedded watermark messages')
parser.add_argument('--decode', action='store_true', default=False,
                    help='if use decoder to extarct watermark messages')

def generate_random_string(length=64):  
    letters_and_digits = string.ascii_letters + string.digits  
    return ''.join(random.choice(letters_and_digits) for i in range(length))

def generate_random_watermark():
    message = generate_random_string()
    assert len(message) == 64
    timestamp = str(time.time()).split('.')[0]
    # print(timestamp)
    assert len(timestamp) == 10
    random_digit = ''.join(random.choice(string.digits) for i in range(4))
    # print(random_digit)
    assert len(random_digit) == 4

    return message+timestamp+random_digit

def dump_watermark(message, file_path):
    with open(file_path, 'w') as f:
        f.writelines(message)

##res = generate_random_watermark()

class Water_Marker:

    def __init__(self, args, seed=0):
        assert os.path.exists(args.checkpoint)
        assert os.path.exists(args.input_videos_dir)
        assert os.path.exists(args.watermark_path)
        ### stable the random seed
        random.seed(seed)

        ### parameters define
        self.mark_index = 0 #459
        self.video_idx  = 0
        self.videomark_label = {}
        self.final_decoded_watermark = {}

        ### accuracy parameters define
        self.correct_nums = 0
        self.total_nums = 0
        self.no_watermark_nums = 0

        ### parameters assignment
        self.checkpoint = args.checkpoint
        self.input_videos_dir = args.input_videos_dir
        self.output_dir = args.output_dir
        self.watermark_path = args.watermark_path
        self.label_path = args.label_path
        self.threshold = args.threshold
        self.use_dictfile = args.use_dictfile
        self.no_accuracy = args.no_accuracy
        self.encode = args.encode
        self.decode = args.decode

        if self.decode:
            assert os.path.exists(self.label_path)

        if not os.path.exists(self.output_dir):
            os.makedirs(self.output_dir)
            
        self.mark_index_msg_dict_path = os.path.join(self.output_dir,"mark_index.pkl")
        print(f"DEBUG: index_msg_dict file path: {self.mark_index_msg_dict_path}")
        self.mark_index_msg_dict = dict()
        if os.path.exists(self.mark_index_msg_dict_path):
            with open(self.mark_index_msg_dict_path,'rb') as f:
                self.mark_index_msg_dict = pickle.load(f)
            
            print("DEBUG: mark index dict detected: \n",self.mark_index_msg_dict)
            self.mark_index = len(self.mark_index_msg_dict)

        ### load model and checkpoints
        self.model = RivaGAN.load(self.checkpoint)
        #print(self.model)
        
        ### load original watermark messages
        self.water_marking = np.array(pickle.load(open(self.watermark_path, 'rb')))
        # self.water_marking = np.load("data/water_marking/wm_10000_100M.npy")
        print('--- water_marking_shape:', self.water_marking.shape)

    def dict2json(self, input_dict):
        json_filepath = self.mark_index_msg_dict_path.replace('.pkl','.json')
        with open(json_filepath, 'w') as f:
            json.dump(json.dumps(input_dict), f)
        print("DEBUG: json file saved: ",json_filepath)
        
    def generate_decoded_watermarks(self, decoded_messages):    
        final_decoded_watermark = {}
        for key, value in decoded_messages.items():
            value = np.array(value)
            biterr_index_list = []
            biterr_value_list = []
            for i in range(value.shape[0]):
                expand_value = value[i].repeat(self.water_marking.shape[0], axis = 0)
                bitwise_avg_err = np.sum(np.abs(expand_value-self.water_marking), \
                                         axis=1) / self.water_marking.shape[1]
            
                ### find min biterr between decoded message and water_marking in library(32bit-watermarks)
                currect_bitwise_err = np.min(bitwise_avg_err)
                currect_bitwise_err_index = np.argmin(bitwise_avg_err)
                biterr_value_list.append(currect_bitwise_err)
                biterr_index_list.append(currect_bitwise_err_index)
        
            biterr_index_array = np.array(biterr_index_list)
            biterr_value_array = np.array(biterr_value_list)

            ### find the most frequent watermark messages
            max_counts_watermark = np.argmax(np.bincount(biterr_index_array))
            
            ### max counts >= value.shape[0]//2
            count_res = np.bincount(biterr_index_array).tolist()
            show_res = {}
            for k,v in enumerate(count_res):
                if v != 0:
                    show_res[k] = v
            print("DEBUG: value.shape: ", value.shape)
            print("DEBUG: decoded watermark ID count: ", show_res)
            print("DEBUG: max_counts_watermark: ", max_counts_watermark)
            if show_res[max_counts_watermark] < value.shape[0]//2:
                print("number of max_counts_watermark less than value.shape[0]//2, set result to no_mark")
                max_counts_watermark = "no_mark"

            ### determine whether the video is embedded with watermark by threshold
            if min(biterr_value_array) < self.threshold:
                self.final_decoded_watermark[key] = max_counts_watermark
            else:
                self.final_decoded_watermark[key] = 'no_mark'

    def encoder(self):
        ### define encoded videos and videomark label saved path
        video_saved_dir = os.path.join(self.output_dir, 'encoded_videos')
        videomark_label_path = os.path.join(self.output_dir, 'videomark_label.npy')
        if not os.path.exists(video_saved_dir):
            os.mkdir(video_saved_dir)

        for video in sorted(os.listdir(self.input_videos_dir)):
            if video.split('.')[-1] not in ['mp4','txt']:
                print(f"DEBUG: not support file detected: {video}")
                continue
            if video.split('.')[-1] == 'txt':
                continue

            print(f"DEBUG: processing video : {video}")
            #### define video input and output paths
            video_input_path = os.path.join(self.input_videos_dir, video)
            #mark_filepath=video_input_path.rsplit('.',maxsplit=1)[0]+'.txt'
            mark_filepath=video_input_path+'.txt'
            if not os.path.exists(mark_filepath):
                print(f"DEBUG: watermark file not exists {mark_filepath}, generate random.")
                res = generate_random_watermark()
                dump_watermark(res, mark_filepath)

            video_saved_path = os.path.join(video_saved_dir, video)
            
            self.mark_index += 1
            #self.mark_index = random.randint(0, 49)
            self.video_idx  += 1
            messages = self.water_marking[self.mark_index]
            messages = torch.FloatTensor(messages).unsqueeze(0).cuda()

            ### add watermark to videos
            ret = self.model.encode(video_input_path, messages, video_saved_path, self.video_idx)
            if ret is not None:
                continue
                
            self.videomark_label[video] = self.mark_index
            with open(mark_filepath,'r') as f:
                watermark_text=f.readlines()[0].strip("\n").strip("\r")
            if len(watermark_text)==0:print("DEBUG: watermark msg text is empty!!")

            self.mark_index_msg_dict[str(self.mark_index)]=watermark_text
            with open(self.mark_index_msg_dict_path,'wb') as f:
                pickle.dump(self.mark_index_msg_dict, f)
            print("DEBUG: saved mark index msg :\n",self.mark_index_msg_dict)
            
            ### to json file
            mark_index_msg_json = self.mark_index_msg_dict
            self.dict2json(mark_index_msg_json)
        
        ### saved videomark label
        np.save(videomark_label_path, self.videomark_label)
        print(f"=> map is {videomark_label_path}")

    def decoder(self):
        ### load decoded watermark messages, if it exists.
        decoded_watermark_path = os.path.join(self.output_dir, 'decoded_watermarks.npy')
        decoded_watermark_dict = {}
        if self.use_dictfile and os.path.exists(decoded_watermark_path):
            decoded_watermark_dict = np.load(decoded_watermark_path, allow_pickle = True)
            decoded_watermark_dict = decoded_watermark_dict.tolist()
            print("=> Decoded watermark dictionary load successfully!!!")
        
        ### else, generate and saved decoded watermark messages
        else:
            print("=> Decoded watermark dictionary not exists, start generating...")
            for video in sorted(os.listdir(self.input_videos_dir)):
                self.video_idx += 1
                video_path = os.path.join(self.input_videos_dir, video)

                ### extract watermark from videos
                decoded_watermarks = self.model.decode(video_path, self.video_idx, self.water_marking)
                decoded_watermark_dict[video] = decoded_watermarks
            
            ### decoded watermark messages saved
            np.save(decoded_watermark_path, decoded_watermark_dict)
        
        ### generate finally videos decoded watermark messages table
        self.generate_decoded_watermarks(decoded_watermark_dict)

        ### saved finally videos decoded watermark messages table
        result_file = os.path.join(self.output_dir, 'final_decoded_watermark.txt')
        ff = open(result_file, 'w')
        for k, v in self.final_decoded_watermark.items():
            ff.write(str(k) + ' ' + str(v) + '\n')
            if str(v) in self.mark_index_msg_dict:
                print(f"{str(v)}:{self.mark_index_msg_dict[str(v)]}")
            else:
                print(f"WARN: mark index not record : {str(v)}")
        ff.close()
        print('=> Final decoded watermark messages saved to {}'.format(result_file))
        print("DEBUG: decoded watermark msg: \n",self.final_decoded_watermark)
        print("DEBUG: saved mark index msg :\n",self.mark_index_msg_dict)

        ### compute decoded accuracy
        if not self.no_accuracy:
            self.decode_accuracy()

    def decode_accuracy(self):
        ### load videomark label
        if os.path.exists(self.label_path):
            videomark_label = np.load(self.label_path, allow_pickle = True)
        else:
            print('=> videomark label file not exists!!!')
        
        videomark_label = videomark_label.tolist()
        print('-' * 80)
        print("DEBUG: input video dir: ", self.input_videos_dir, '\n',self.final_decoded_watermark)
        print("DEBUG: videomark_label: ", videomark_label)
        for video_name in os.listdir(self.input_videos_dir):
            if video_name not in self.final_decoded_watermark:
                print("ERROR: not in self.final_decoded_watermark: ", video_name)
            if video_name not in videomark_label:
                print(f"INFO: video {video_name} not in videomark_label record")
                continue

            self.total_nums += 1
            if self.final_decoded_watermark[video_name] == 'no_mark':
                self.no_watermark_nums += 1
                print(f'=> Video: [{video_name}] no watermark messasges!')
            elif self.final_decoded_watermark[video_name] == videomark_label[video_name]:
                self.correct_nums += 1
            else:
                print(f'=> Video: [{video_name}] has watermark messages, but decode error!')
        
        ### print informations(accuracy, decode correct nums, no watermark nums)
        if self.total_nums == self.no_watermark_nums:
            print("=> All videos have no watermark messages!!!")
        else:
            accuracy = self.correct_nums / (self.total_nums - self.no_watermark_nums)
            print("=> No watermark messages video nums: {}".format(self.no_watermark_nums))
            print("=> Decode correct_nums/has_watermark_nums: [{}/{}], accuracy: {}%".format(
                                                              self.correct_nums,
                                                              self.total_nums - self.no_watermark_nums,
                                                              accuracy*100.0))

   
if __name__ == "__main__":
    
    args = parser.parse_args()
    water_marker = Water_Marker(args)
    
    print("=" * 80)

    if args.encode:
        print('=> Run Encoder ...')
        water_marker.encoder()
    elif args.decode:
        print('=> Run Decoder ...')
        water_marker.decoder()
     
    print("=" * 80)
